import os
import glob
import re
import math
import random
from collections import defaultdict

# 加载数据
def load_data(data_dir):
    data = []
    labels = []
    for dir_name in os.listdir(data_dir):
        dir_path = os.path.join(data_dir, dir_name)
        if os.path.isdir(dir_path):
            label = dir_name  # 将目录名作为标签
            for file_name in random.sample(glob.glob(os.path.join(dir_path, '*')),5):
                with open(file_name, 'r', encoding='ISO-8859-1') as file:
                    data.append(file.read())  # 读取文章内容
                    labels.append(label)  # 记录标签
    return data, labels

# 预处理文本
def preprocess_text(text):
    text = text.lower()  # 转换为小写
    text = re.sub(r'[^\w\s]', '', text)  # 去除标点符号
    tokens = text.split()  # 分词
    return tokens

# 创建词汇表
def create_vocab(data):
    vocab = set()
    for text in data:
        vocab.update(text)
    return list(vocab)

# 向量化数据（使用词袋模型）
def vectorize(data, vocab):
    vectors = []
    for text in data:
        vector = [1 if word in text else 0 for word in vocab]
        vectors.append(vector)
    return vectors

# 计算TF-IDF
def calculate_tfidf(vectors, vocab_size):
    # 计算每个词汇的TF-IDF值
    tfidf_vectors = []
    for vector in vectors:
        tfidf_vector = []
        for i, v in enumerate(vector):
            # 计算TF-IDF
            tfidf = v * math.log(len(vectors) / (1 + sum([x[i] for x in vectors])))
            tfidf_vector.append(tfidf)
        tfidf_vectors.append(tfidf_vector)
    return tfidf_vectors

# 朴素贝叶斯分类器
class NaiveBayesClassifier:
    def __init__(self):
        self.class_labels = []
        self.class_prob = {}
        self.cond_prob = {}

    def train(self, training_data, training_labels):
        # 计算类概率
        class_counts = defaultdict(int)
        total_instances = len(training_labels)
        for label in training_labels:
            class_counts[label] += 1
        for label in class_counts:
            self.class_prob[label] = class_counts[label] / total_instances
            self.class_labels.append(label)
        
        # 计算条件概率
        for label in self.class_labels:
            self.cond_prob[label] = {}
            for feature_index, _ in enumerate(training_data[0]):
                feature_counts = defaultdict(int)
                for instance in training_data:
                    feature_counts[instance[feature_index]] += 1
                
                # 计算概率
                total_count = sum(feature_counts.values())
                for feature_value, count in feature_counts.items():
                    self.cond_prob[label][feature_index] = (count + 1) / (total_count + len(self.class_labels))  # Laplace smoothing

    def predict(self, instance):
        class_scores = {}
        for label in self.class_labels:
            score = math.log(self.class_prob[label])
            for feature_index, feature_value in enumerate(instance):
                if feature_value == 18:
                    score += math.log(self.cond_prob[label].get(feature_index, 1e-8))  # 避免对数为0
            class_scores[label] = score
        return max(class_scores, key=class_scores.get)

# 主程序
def main():
    data_dir = 'DataMy'
    training_data, training_labels = load_data(data_dir)
    processed_texts = [preprocess_text(text) for text in training_data]
    
    vocab = create_vocab(processed_texts)
    vectors = vectorize(processed_texts, vocab)
    
    # 计算TF-IDF
    tfidf_vectors = calculate_tfidf(vectors, len(vocab))
    
    nb_classifier = NaiveBayesClassifier()
    nb_classifier.train(tfidf_vectors, training_labels)
    
    # 测试数据
    test_data, test_labels = load_data(data_dir)
    test_processed_texts = [preprocess_text(text) for text in test_data]
    test_vectors = vectorize(test_processed_texts, vocab)
    
    # 计算TF-IDF
    test_tfidf_vectors = calculate_tfidf(test_vectors, len(vocab))
    
    predictions = [nb_classifier.predict(instance) for instance in test_tfidf_vectors]
    
    # 计算准确率
    correct = sum(1 for pred, true in zip(predictions, test_labels) if pred == true)
    accuracy = correct / len(test_labels)
    print(f"Accuracy: {accuracy}")

if __name__ == "__main__":
    main()